﻿// Console.cpp : 此文件包含 "main" 函数。程序执行将在此处开始并结束。
//

#include "pch.h"
#include <iostream>
#include <shlwapi.h>
#include <io.h>
#include <direct.h>

#pragma comment(lib, "shlwapi.lib")
#pragma comment( lib, "version.lib" )

#include "Capture.h"

#include <GdiPlus.h>
#pragma  comment(lib, "gdiplus.lib")
#pragma  comment(lib, "Gdi32.lib")


#define Client_MapData	"Client_MapData"
#include <tlhelp32.h>



HWND g_ClientHwnd;
enum CAPTURE_FORMAT
{
	CF_BMP,
	CF_JPG,
	CF_PNG
};

static int GetEncoderClsid(const WCHAR* format, CLSID* pClsid)
{
	UINT num = 0;                     // number of image encoders     
	UINT size = 0;                   // size of the image encoder array in bytes     
	Gdiplus::ImageCodecInfo* pImageCodecInfo = NULL;
	Gdiplus::GetImageEncodersSize(&num, &size);
	if (size == 0)
		return -1;     //   Failure     

	pImageCodecInfo = (Gdiplus::ImageCodecInfo*)(malloc(size));
	if (pImageCodecInfo == NULL)
		return -1;     //   Failure     

	Gdiplus::GetImageEncoders(num, size, pImageCodecInfo);
	for (UINT j = 0; j < num; ++j)
	{
		if (wcscmp(pImageCodecInfo[j].MimeType, format) == 0)
		{
			*pClsid = pImageCodecInfo[j].Clsid;
			free(pImageCodecInfo);
			return j;     //   Success     
		}
	}
	free(pImageCodecInfo);
	return -1;     //   Failure     
}

int SendDataMessage(HWND hw, std::string msg)
{
	if (!hw) { hw = g_ClientHwnd; } if (!hw) return 1;
	printf("\nSDM:%x\n", g_ClientHwnd);
	COPYDATASTRUCT data = { 0 };
	data.cbData = msg.length() + 1;
	data.lpData = (void*)msg.data();
	::SendMessageA(hw, WM_COPYDATA, 0, (LPARAM)&data);
	return 0;
}


static int ReadMMem(const char *name, void* buffer, int size)
{
	HANDLE hMap = OpenFileMappingA(FILE_MAP_READ | FILE_MAP_WRITE, FALSE, name);
	if (hMap == INVALID_HANDLE_VALUE)
	{
		printf("error at read mmem %d\n", GetLastError());
		return 1;
	}
	void* pMapView = (void*)MapViewOfFile(hMap, FILE_MAP_READ | FILE_MAP_WRITE, 0, 0, 0);

	if (!pMapView)
	{
		printf("error at rmap view %d\n", GetLastError());
		CloseHandle(hMap);
		return 2;
	}

	memcpy(buffer, pMapView, size);

	UnmapViewOfFile(pMapView);
	CloseHandle(hMap);
	return 0;
}
static void xyDebugWindowBind(const char* title)
{
	if (FindWindowA(0, title))
	{
		DWORD pid = 0;
		GetWindowThreadProcessId(FindWindowA(0, title), &pid);
		AttachConsole(pid);
		freopen("CONIN$", "r", stdin);
		freopen("CONOUT$", "w", stdout);
		freopen("CONOUT$", "w", stderr);
	}
}
static const char* GetAppPath()
{
	static char fpath[256];
	if (fpath[0] == NULL)
	{
		GetModuleFileNameA(NULL, fpath, 256);
		PathRemoveFileSpecA(fpath);
	}
	return fpath;
}
static wchar_t* GetAppPathW()
{
	static wchar_t fpath[256];
	if (fpath[0] == NULL)
	{
		GetModuleFileNameW(NULL, fpath, 256);
		PathRemoveFileSpecW(fpath);
	}
	return fpath;
}

Gdiplus::Bitmap* getBitmap(int x, int y, int w, int h, bool bSave) {
	HWND hWnd = ::GetDesktopWindow();
	HDC hDC = ::GetDC(hWnd);
	RECT rect;
	::GetClientRect(hWnd, &rect);
	//HDC拷贝至内存DC;
	HDC memDC = CreateCompatibleDC(hDC);
	HBITMAP hBmp = CreateCompatibleBitmap(hDC, w, h);
	::SelectObject(memDC, hBmp);
	::BitBlt(memDC, 0, 0, w, h, hDC, x, y, SRCCOPY | CAPTUREBLT);

	//保存图片;
	Gdiplus::Bitmap* bmWnd = Gdiplus::Bitmap::FromHBITMAP(hBmp, NULL);

	if (bSave) {
		std::wstring file_bmp = GetAppPathW();
		file_bmp += L"\\capture.bmp";
		CLSID Clsid;
		GetEncoderClsid(L"image/bmp", &Clsid);
		bmWnd->Save(file_bmp.c_str(), &Clsid);
	}
		

	::SelectObject(memDC, (HBITMAP)NULL);
	::DeleteDC(memDC);
	::DeleteObject(hBmp);
	::ReleaseDC(hWnd, hDC);

	return bmWnd;
}

bool	bmpCompare(Gdiplus::Bitmap* bmWnd, Gdiplus::Bitmap* wndWnd) {
	UINT bmpHeight = bmWnd->GetHeight();
	UINT bmpWidth = bmWnd->GetWidth();
	UINT wndHeight = wndWnd->GetHeight();
	UINT wndWidth = wndWnd->GetWidth();
	Gdiplus::Color bmColor;
	Gdiplus::Color wndColor;
	for (size_t x = 0; x < bmpWidth; x++)
		for (size_t y = 0; y < bmpHeight; y++)
		{
			bmWnd->GetPixel(x, y, &bmColor);
			wndWnd->GetPixel(x, y, &wndColor);
			if (bmColor.GetValue() != wndColor.GetValue()) {
				return false;
			}
		}
	return true;
}

int capture_start(int x, int y, int w, int h) {
	if (x <= 0 || y <= 0 || w <= 0 || h <= 0) {
		printf("error at capture_start");
		return 1;
	}
	
	Gdiplus::GdiplusStartupInput gdiplusStartupInput;
	ULONG_PTR gdiplusToken;
	Gdiplus::GdiplusStartup(&gdiplusToken, &gdiplusStartupInput, NULL);

	Gdiplus::Bitmap* bmWnd = getBitmap(x, y, w, h, false);

	for (;;){
		Sleep(200);
		Gdiplus::Bitmap* wndWnd = getBitmap(x, y, w, h, false);
		if (!bmpCompare(bmWnd, wndWnd)) {
			printf("break\n");
			getBitmap(x, y, w, h, true);

			std::string file_bmp = GetAppPath();
			file_bmp += "\\capture.bmp";

			for (size_t i = 0; i < file_bmp.length(); i++)
			{
				if (file_bmp[i] == '\\') file_bmp[i] = '/';
			}
			printf("file:%s\n", file_bmp.c_str());
			std::string msg = "msg_capture_notify('";
			msg += file_bmp;
			msg += "')";
			printf("%s", msg.c_str());
			SendDataMessage(0, msg);
			break;
		}
		delete wndWnd;
	}
	delete bmWnd;
	Gdiplus::GdiplusShutdown(gdiplusToken);
	return 0;
}

bool boot_start_mutex()
{
	HANDLE m_hMutex = ::CreateMutexA(NULL, TRUE, "syxzsCapture");
	if (GetLastError() == ERROR_ALREADY_EXISTS) //程序已经运行
	{
		m_hMutex = NULL;
		return true;
	}
	return false;
}


void findWnd(DWORD pid) {
	DWORD id = 0;
	char title[255];
	char name[255];
	HWND hwnd = ::FindWindowA(NULL, NULL);
	while (hwnd != NULL) {
		::GetWindowTextA(hwnd, title, MAX_PATH);
		::GetClassNameA(hwnd, name, MAX_PATH);
		if (GetWindowThreadProcessId(hwnd, &id) && pid == id) {
			printf("%X:%s[%s]\n", hwnd, title, name);
		}
		hwnd = ::FindWindowExA(0, hwnd, NULL, NULL);
	}
}

int globlePid = 0;
HHOOK keyHook = NULL;
HHOOK mouseHook = NULL;

//键盘钩子过程
LRESULT CALLBACK keyProc(int nCode, WPARAM wParam, LPARAM lParam)
{
	//在WH_KEYBOARD_LL模式下lParam 是指向KBDLLHOOKSTRUCT类型地址
	KBDLLHOOKSTRUCT *pkbhs = (KBDLLHOOKSTRUCT *)lParam;
	//如果nCode等于HC_ACTION则处理该消息，如果小于0，则钩子子程就必须将该消息传递给 CallNextHookEx
	//if (nCode == HC_ACTION){
	if (pkbhs->vkCode == VK_ESCAPE && GetAsyncKeyState(VK_CONTROL) & 0x8000 && GetAsyncKeyState(VK_SHIFT) & 0x8000) {
		std::cout << "Ctrl+Shift+Esc";
		return 1;
	}
	else if (pkbhs->vkCode == VK_ESCAPE && GetAsyncKeyState(VK_CONTROL) & 0x8000) {
		std::cout << "Ctrl+Esc";
		return 1;
	}
	else if (pkbhs->vkCode == VK_TAB && pkbhs->flags & LLKHF_ALTDOWN) {
		std::cout << "Alt+Tab";
		return 1;
	}
	else if (pkbhs->vkCode == VK_ESCAPE && pkbhs->flags &LLKHF_ALTDOWN) {
		std::cout << "Alt+Esc";
		return 1;
	}
	else if (pkbhs->vkCode == VK_LWIN || pkbhs->vkCode == VK_RWIN) {
		std::cout << "LWIN/RWIN";
		return 1;
	}
	else if (pkbhs->vkCode == VK_F4 && pkbhs->flags & LLKHF_ALTDOWN) {
		std::cout << "Alt+F4";
		return 1;
	}

	//return 1;//返回1表示截取消息不再传递,返回0表示不作处理,消息继续传递
	//}
	return CallNextHookEx(keyHook, nCode, wParam, lParam);
}

//鼠标钩子过程
LRESULT CALLBACK mouseProc(int nCode, WPARAM wParam, LPARAM lParam)
{
	return 1;
}

//卸载钩子
void unHook()
{
	UnhookWindowsHookEx(keyHook);
	//  UnhookWindowsHookEx(mouseHook);
	printf("unHook\n");
}

//安装钩子,调用该函数即安装钩子
void setHook()
{
	//这两个底层钩子,不要DLL就可以全局
	//底层键盘钩子
	keyHook = SetWindowsHookExA(WH_KEYBOARD_LL, keyProc, GetModuleHandleA(0), 0);
	printf("keyHook:%d\n", keyHook);
	//底层鼠标钩子
	//    mouseHook =SetWindowsHookEx( WH_MOUSE_LL,mouseProc,GetModuleHandle(NULL),0);
}
//提升权限
void EnableDebugPriv()
{
	HANDLE hToken;
	LUID sedebugnameValue;
	TOKEN_PRIVILEGES tkp;

	if (!OpenProcessToken(GetCurrentProcess(),
		TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
		return;
	}
	if (!LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &sedebugnameValue)) {
		CloseHandle(hToken);
		return;
	}
	tkp.PrivilegeCount = 1;
	tkp.Privileges[0].Luid = sedebugnameValue;
	tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
	if (!AdjustTokenPrivileges(hToken, FALSE, &tkp, sizeof(tkp), NULL, NULL)) {
		CloseHandle(hToken);
		return;
	}
}


char* ConvertLPWSTRToLPSTR(LPWSTR lpwszStrIn)
{
	LPSTR pszOut = NULL;
	if (lpwszStrIn != NULL)
	{
		int nInputStrLen = wcslen(lpwszStrIn);

		// Double NULL Termination  
		int nOutputStrLen = WideCharToMultiByte(CP_ACP, 0, lpwszStrIn, nInputStrLen, NULL, 0, 0, 0) + 2;
		pszOut = new char[nOutputStrLen];

		if (pszOut)
		{
			memset(pszOut, 0x00, nOutputStrLen);
			WideCharToMultiByte(CP_ACP, 0, lpwszStrIn, nInputStrLen, pszOut, nOutputStrLen, 0, 0);
		}
	}
	return pszOut;
}

//冻结
void Freeze()
{
	//枚举进程信息
	PROCESSENTRY32 pe32;
	pe32.dwSize = sizeof(pe32);
	HANDLE hProcessSnap = ::CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

	int processPid;
	//CString strTmp;
	BOOL b = ::Process32First(hProcessSnap, &pe32);
	while (b)
	{
		processPid = pe32.th32ProcessID;
		char *exeFile = ConvertLPWSTRToLPSTR(pe32.szExeFile);
		if (strcmp(exeFile, "winlogon.exe") == 0)
		{
			break;
		}

		delete[] exeFile;

		b = ::Process32Next(hProcessSnap, &pe32);
	}
	::CloseHandle(hProcessSnap);

	THREADENTRY32 th32;
	th32.dwSize = sizeof(th32);
	HANDLE hThreadSnap = ::CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
	globlePid = processPid;
	unsigned long Pid;
	Pid = processPid;
	b = ::Thread32First(hThreadSnap, &th32);
	while (b)
	{
		if (th32.th32OwnerProcessID == Pid)
		{
			HANDLE oth = OpenThread(THREAD_ALL_ACCESS, FALSE, th32.th32ThreadID);
			if (!(::SuspendThread(oth)))
			{
				std::cout << "freeze successed";
			}
			else
			{
				std::cout << "freeze failed";
			}
			CloseHandle(oth);
			break;
		}
		::Thread32Next(hThreadSnap, &th32);
	}
	::CloseHandle(hThreadSnap);
}

void unFreeze()
{
	unsigned long Pid;
	Pid = globlePid;

	THREADENTRY32 th32;
	th32.dwSize = sizeof(th32);

	HANDLE hThreadSnap = ::CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
	BOOL b = ::Thread32First(hThreadSnap, &th32);
	while (b)
	{
		if (th32.th32OwnerProcessID == Pid)
		{
			HANDLE oth = OpenThread(THREAD_ALL_ACCESS, FALSE, th32.th32ThreadID);
			if (::ResumeThread(oth))
			{
				std::cout << "unfreeze successed";
			}
			else
			{
				std::cout << "unfreeze failed";
			}
			CloseHandle(oth);
			break;
		}
		::Thread32Next(hThreadSnap, &th32);
	}
	::CloseHandle(hThreadSnap);
}

bool checkProc(std::wstring name)
{
	bool res = false;
	//枚举进程信息
	PROCESSENTRY32 pe32;
	pe32.dwSize = sizeof(pe32);
	HANDLE hProcessSnap = ::CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

	BOOL b = ::Process32First(hProcessSnap, &pe32);
	while (b)
	{
		if (pe32.szExeFile == name)
		{
			res = true; break;
		}
		b = ::Process32Next(hProcessSnap, &pe32);
	}
	::CloseHandle(hProcessSnap);
	return res;
}

// 重启计算机
void on_msg_shutdown_system()
{
	HANDLE           hToken = NULL;
	TOKEN_PRIVILEGES tkp = { 0 };

	typedef enum _SHUTDOWN_ACTION
	{
		ShutdownNoReboot,
		ShutdownReboot,
		ShutdownPowerOff
	} SHUTDOWN_ACTION;

	typedef DWORD(WINAPI* lpNtShutdownSystem)(SHUTDOWN_ACTION Action);
	lpNtShutdownSystem NtShutdownSystem = NULL;

	OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, &hToken);
	LookupPrivilegeValue(NULL, SE_SHUTDOWN_NAME, &tkp.Privileges[0].Luid);
	tkp.PrivilegeCount = 1;
	tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
	AdjustTokenPrivileges(hToken, FALSE, &tkp, 0, NULL, 0);
	CloseHandle(hToken);

	HMODULE hNtdll = GetModuleHandleA("ntdll.dll");
	if (hNtdll)
	{
		NtShutdownSystem = (lpNtShutdownSystem)GetProcAddress(hNtdll, "NtShutdownSystem");
		if (NtShutdownSystem)
		{
			NtShutdownSystem(ShutdownReboot);
		}
	}
}

void recovery() {
	static int only;
	if (FindWindowA("GlassWndClass-GlassWindowClass-3", "main")) {
		if (!only) {
			unHook();
			unFreeze();
			only = 1;
		}
	}
}


void check() {
	Sleep(4000);
	for (;;){
		recovery();
		if (!checkProc(L"java.exe")) {
			on_msg_shutdown_system();
			break;
		}
		Sleep(1000);
	}
}

void checkWnd() {
	size_t length = 60 * 5;
	for (size_t i = 0; i < length; i++)
	{
		Sleep(1000);
		if (FindWindowA("GlassWndClass-GlassWindowClass-3", "main")) {
			return;
		}

		HWND hwnd = FindWindowA(0, "网竞教育");
		if (hwnd) {
			SetWindowPos(hwnd, HWND_TOPMOST, 0, 0, 0, 0, SWP_NOMOVE | SWP_NOSIZE | SWP_SHOWWINDOW);
		}
	}
	on_msg_shutdown_system();
}

void startCheck() {
	DWORD dwThreadID = 0;
	CloseHandle(CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)check, 0, 0, &dwThreadID));
	CloseHandle(CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)checkWnd, 0, 0, &dwThreadID));
}

int yymain(int argc, const char* args[]) {
	checkWnd();
}

int main(int argc, const char* args[])
{
	xyDebugWindowBind("管理员:  wy-capture");
	SCREENSEL  scrSelect = { 0 };
	LPCWSTR	   szWndClassName = L"Capture";
	MSG        msg;
	HDC        hdcRealScreen;
	HINSTANCE hInstance = 0;
	if (boot_start_mutex()){
		printf("syxzsCapture  runing");
		return 0;
	}

	if (argc != 2) {
		return 0;
	}

	// Setup our selection parameters
	SetRectEmpty(&scrSelect.rcSelect);
	SetRect(&scrSelect.rcScreen,
		GetSystemMetrics(SM_XVIRTUALSCREEN), GetSystemMetrics(SM_YVIRTUALSCREEN),
		GetSystemMetrics(SM_CXVIRTUALSCREEN), GetSystemMetrics(SM_CYVIRTUALSCREEN));

	// Create cursor from raw data
	scrSelect.hcrCross = CreateCursor(hInstance, 12, 21, 32, 32, bCrossCursor, bCrossCursor + 128);

	// create pen for selection border
	scrSelect.hpnBorder = CreatePen(PS_DOT, 1, RGB(255, 255, 255));

	// Get screen DC
	hdcRealScreen = GetDC(NULL);
	scrSelect.hbmScreen = CreateCompatibleBitmap(hdcRealScreen,
		scrSelect.rcScreen.right, scrSelect.rcScreen.bottom);
	if (!scrSelect.hbmScreen)
	{
		goto cleanup;
	}
	scrSelect.hdcScreen = CreateCompatibleDC(hdcRealScreen);
	if (!scrSelect.hdcScreen)
	{
		goto cleanup;
	}
	scrSelect.ndcScreen = SaveDC(scrSelect.hdcScreen);
	scrSelect.hbmMem = CreateCompatibleBitmap(hdcRealScreen,
		scrSelect.rcScreen.right, scrSelect.rcScreen.bottom);
	if (!scrSelect.hbmMem)
	{
		goto cleanup;
	}
	scrSelect.hdcMem = CreateCompatibleDC(hdcRealScreen);
	if (!scrSelect.hdcMem)
	{
		goto cleanup;
	}
	scrSelect.ndcMem = SaveDC(scrSelect.hdcScreen);

	SelectObject(scrSelect.hdcScreen, scrSelect.hbmScreen);
	SelectObject(scrSelect.hdcMem, scrSelect.hbmMem);
	BitBlt(scrSelect.hdcScreen, 0, 0, scrSelect.rcScreen.right,
		scrSelect.rcScreen.bottom, hdcRealScreen, 0, 0, SRCCOPY);

	ReleaseDC(NULL, hdcRealScreen);
	hdcRealScreen = NULL;

	if (!InitRegisterClass(hInstance, szWndClassName, &scrSelect))
	{
		goto cleanup;
	}

	if (InitInstance(hInstance, szWndClassName, SW_SHOW, &scrSelect) == FALSE)
	{
		goto cleanup;
	}

	
	char command[1024];
	sprintf_s(command, sizeof(command), "/c java -jar \"./client-app.jar\" %s", args[1]);
	ShellExecuteA(0, "open", "cmd.exe", command, 0, SW_HIDE);

	EnableDebugPriv();
	Freeze();
	printf("setHook\n");
	setHook();

	startCheck();
	while (GetMessage(&msg, NULL, 0, 0))
	{
		if (msg.message == WM_USER) break;
		TranslateMessage(&msg);
		DispatchMessage(&msg);
	}
	
	
cleanup:
	if (hdcRealScreen)
		ReleaseDC(NULL, hdcRealScreen);
	if (scrSelect.hdcMem) {
		RestoreDC(scrSelect.hdcMem, scrSelect.ndcMem);
		DeleteDC(scrSelect.hdcMem);
	}
	if (scrSelect.hbmMem)
		DeleteObject(scrSelect.hbmMem);
	if (scrSelect.hdcScreen) {
		RestoreDC(scrSelect.hdcScreen, scrSelect.ndcScreen);
		DeleteDC(scrSelect.hdcScreen);
	}
	if (scrSelect.hpnBorder)
		DeleteObject(scrSelect.hpnBorder);
	if (scrSelect.hcrCross)
		DestroyCursor(scrSelect.hcrCross);

	UnregisterClass(szWndClassName, hInstance); // unregister wcexSelect

	return (int)msg.wParam;
}



// 运行程序: Ctrl + F5 或调试 >“开始执行(不调试)”菜单
// 调试程序: F5 或调试 >“开始调试”菜单

// 入门提示: 
//   1. 使用解决方案资源管理器窗口添加/管理文件
//   2. 使用团队资源管理器窗口连接到源代码管理
//   3. 使用输出窗口查看生成输出和其他消息
//   4. 使用错误列表窗口查看错误
//   5. 转到“项目”>“添加新项”以创建新的代码文件，或转到“项目”>“添加现有项”以将现有代码文件添加到项目
//   6. 将来，若要再次打开此项目，请转到“文件”>“打开”>“项目”并选择 .sln 文件
